import torch
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from seg_dataset import mydata

from torch.utils.data import DataLoader
from I_FCN import VGG16
import torch




#


model=VGG16().to('cuda')



PATH = 'F:\\lunwen\\save_weights\\GoogLeNet_backbone_224_aspp_nonormal9.pth'
model.load_state_dict(torch.load(PATH)['model'])


num=6

val_data = mydata('C:\\Users\\HDH\\Desktop\\data2', 224, 'val')


val_loader = DataLoader(val_data, batch_size=6,shuffle=True)
image, mask = next(iter(val_loader))
pred_mask = model(image)


for i in range(num):
    plt.figure(figsize=(30, 30))
    plt.subplot(num, 3, i*3+1)
    plt.imshow(image[i].permute(1,2,0).cpu().numpy())
    plt.subplot(num, 3, i*3+2)
    plt.imshow(mask[i].cpu().numpy())
    plt.subplot(num, 3, i*3+3)
    plt.imshow(torch.argmax(pred_mask[i].permute(1,2,0), axis=-1).detach().numpy())
    plt.show()
